#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Created on Mon Feb  5 14:20:02 2018

@author: annabelle
"""
import pickle
import os
import numpy as np
import matplotlib.pyplot as plt
from stringCore import new_string


def generateSnapshots(case, lengthDist, filter, dopings):
    di = 'V:\\Paper Data\AFM_StringTheory\Snapshots_experiment\\'
    do = 'V:\\Paper Data\AFM_StringTheory\\'

    for fn in os.listdir(di):
        # only analyze cold or hot grouped data (T < 0.7J or T > 1.3 J), or temperature slice at doping 10 (0, 6, 12)
        if fn.startswith(filter) and 'D0' in fn:
            with open(di+fn, 'rb') as f:
                amses = pickle.load(f)

            print '  Using filename: {} with {}+{} shots'.format(fn, len(amses[0]), len(amses[1])+len(amses[2]))

            for doping in dopings:
                print '    Generating doping {}'.format(doping)
                amsesout = []
                for k, ams in enumerate(amses):
                    amsout = []
                    for i, img in enumerate(ams):
                        img = np.array(img)
                        m = img.copy()

                        # reduce to minimal array around system size
                        roi = np.where(m != 0)
                        # find corners of big roi
                        xMin = min(np.unique(roi[0]))
                        xMax = max(np.unique(roi[0]))
                        yMin = min(np.unique(roi[1]))
                        yMax = max(np.unique(roi[1]))

                        # leave a 1-site buffer all around
                        m = m[xMin-1:xMax + 2, yMin-1:yMax + 2]
                        roi = np.abs(m)

                        # set number of sites to determine number of strings - include number fluctuations with random
                        totalSites = np.sum(roi)  # mX*mY
                        # numStrings = 0
                        # for i in range(totalSites):
                        #     r = random.random()
                        #     if r < 0.01 * doping:
                        #         numStrings += 1
                        numStrings = int(doping/100.0 * totalSites + 0.5)

                        # put strings into the picture
                        holeEndingPos = []
                        stringLengths = []
                        # if numStrings > 0:
                        #     fig, axes = plt.subplots(1, 3)
                        #     axes[0].imshow(m)
                        for _ in range(numStrings):
                            string = None
                            while string == None:
                                m, holeEndingPos, movesList, string = new_string(m, 0, 0, holeEndingPos, np.cumsum(lengthDist))
                            stringLengths.append(len(string) - 1)
                        #     axes[1].plot(np.transpose(string)[1], np.transpose(string)[0], zorder=10)
                        #     axes[0].plot(np.transpose(string)[1], np.transpose(string)[0], zorder=10)
                        #     axes[1].plot(string[0][1], string[0][0], '.', zorder=11)
                        #     axes[0].plot(string[0][1], string[0][0], '.', zorder=11)
                        # if numStrings > 0:
                        #     axes[1].imshow(m)
                        #     axes[1].set_title(str(numStrings))
                        # set everything outside of roi to 0 again (because strings can run into roi (not out currently)!)
                        m[roi == 0] = 0
                        # if numStrings > 0:
                        #     axes[2].imshow(m), plt.show()

                        amsout.append(np.copy(m))
                    amsesout.append(np.copy(amsout))

                if case == 'peak0':
                    folder = 'Snapshots_sprinkled\\'
                else:
                    folder = 'Snapshots_strings\\'

                if filter.startswith('D'):
                    fo = "{}_D{}_{}.pkl".format(case, doping, fn.split('_')[1][:-4])
                else:
                    fo = "{}_{}_D{}.pkl".format(case, filter, doping)
                with open(do + folder + fo, 'wb') as f:
                    pickle.dump(amsesout, f)
                print "    Saving to {}".format(fo)


if __name__ == '__main__':
    filters = ['cold', 'hot', 'D']

    for filter in filters:
        if filter == 'cold':
            # load analytic string length distribution
            T0 = 0.6
            Thot = 0.65
            Tcold = 0.55
            stringLengths_fn = "V:\Paper Data\AFM_StringTheory\\analyticGST_lenDist_Tdep.pickle"
            with open(stringLengths_fn, 'rb') as f:
                stringLengths = pickle.load(f)
            geometricLengths = stringLengths['stringLengthsData'][np.where(stringLengths['temperatures'] == T0)[0][0]]

            # sprinkling
            sprinklingLengths = np.zeros_like(geometricLengths)
            sprinklingLengths[0] = 1

            # string length delta peak
            peakLengths = []
            peaks = [1, 2, 3, 16]
            for peak in peaks:
                peakLengths.append(np.zeros_like(geometricLengths))
                peakLengths[-1][peak] = 1

            # different participation ratios with gst0.60J
            participLengths = []
            participRatios = [0.8, 0.5]
            for ratio in participRatios:
                participLengths.append(np.copy(geometricLengths))
                participLengths[-1] *= ratio
                participLengths[-1][0] += 1 - ratio

            cases = {'gst0.60J': geometricLengths,
                     'gst0.65J': stringLengths['stringLengthsData'][np.where(stringLengths['temperatures'] == Thot)[0][0]],
                     'gst0.55J': stringLengths['stringLengthsData'][np.where(stringLengths['temperatures'] == Tcold)[0][0]],
                     'peak0': sprinklingLengths,
                     'peak1': peakLengths[0],
                     'peak2': peakLengths[1],
                     'peak3': peakLengths[2],
                     'peakinf': peakLengths[3],
                     'pR0.8': participLengths[0],
                     'pR0.5': participLengths[1]
                     }

            dopings = range(0, 35, 2)

        elif filter == 'hot':
            # load analytic string length distribution
            T0 = 0.8
            stringLengths_fn = "V:\Paper Data\AFM_StringTheory\\analyticGST_lenDist_Tdep.pickle"
            with open(stringLengths_fn, 'rb') as f:
                stringLengths = pickle.load(f)
            geometricLengths = stringLengths['stringLengthsData'][np.where(stringLengths['temperatures'] == T0)[0][0]]

            # sprinkling
            sprinklingLengths = np.zeros_like(geometricLengths)
            sprinklingLengths[0] = 1

            cases = {'gst0.80J': geometricLengths,
                     'peak0': sprinklingLengths,
                     }

            dopings = [0, 10]

        elif filter == 'D':
            # load analytic string length distribution, here it will be unused
            T0 = 0.8
            stringLengths_fn = "V:\Paper Data\AFM_StringTheory\\analyticGST_lenDist_Tdep.pickle"
            with open(stringLengths_fn, 'rb') as f:
                stringLengths = pickle.load(f)
            geometricLengths = stringLengths['stringLengthsData'][np.where(stringLengths['temperatures'] == T0)[0][0]]

            # sprinkling
            sprinklingLengths = np.zeros_like(geometricLengths)
            sprinklingLengths[0] = 1

            cases = {'peak0': sprinklingLengths,
                     }

            dopings = [0, 6, 10, 12]

        for case, dist in iter(cases.items()):
            print "GENERATING ", case, filter
            if 'peak' in case:
                fmt = '.'
            else:
                fmt = '-'
            plt.plot(dist[:17], fmt, label=case)

            generateSnapshots(case, dist, filter, dopings)

    plt.legend()
    plt.show()

